import numpy as np
import mindspore
from mindspore import ops
from mindspore import Tensor, CSRTensor, COOTensor, RowTensor
from mindspore.common.initializer import One, Normal


def show_tensor_attr(tensor:Tensor):

    # tensor 的 形状
    tensor_shape = tensor.shape

    # tensor 的 数据类型
    tensor_dtype = tensor.dtype

    # tensor 单个元素占用的字节数
    tensor_itemsize = tensor.itemsize

    # tensor 占用的字节数
    tensor_nbytes = tensor.nbytes

    # tensor 的秩，不同于矩阵的秩，这里指的是 len(tensor.shape)
    tensor_ndim = tensor.ndim

    # tensor 的 元素的数量
    tensor_size = tensor.size

    # tensor 每一维所需要的字节数
    tensor_strides = tensor.strides

    for name, value in locals().items():
        print(f"{name}: {value}")
    print()


# 根据数据直接生成
def gen_tensor_base():
    data = [1, 0, 1, 0]
    x_data = Tensor(data)
    show_tensor_attr(x_data)

# 从 NumPy 数组生成
def gen_tensor_numpy():
    data = [1, 0, 1, 0]
    np_array = np.array(data)
    x_np = Tensor(np_array)
    show_tensor_attr(x_np)

# 使用 init 初始化器构造张量
def gen_tensor_init():
    tensor1 = mindspore.Tensor(shape=(2, 2), dtype=mindspore.float32, init=One())
    # Initialize a tensor from normal distribution
    tensor2 = mindspore.Tensor(shape=(2, 2), dtype=mindspore.float32, init=Normal())

    print("tensor1:")
    show_tensor_attr(tensor1)

    print("tensor2:")
    show_tensor_attr(tensor2)

# 继承另一个张量的属性，形成新的张量
def gen_tensor_extend():
    data = [1, 0, 1, 0]
    x_data = Tensor(data)

    x_ones = ops.ones_like(x_data)
    print(f"Ones Tensor: \n {x_ones} \n")

    x_zeros = ops.zeros_like(x_data)
    print(f"Zeros Tensor: \n {x_zeros} \n")


def show_tensor_index():
    tensor = Tensor(np.array([[0, 1], [2, 3]]).astype(np.float32))

    print("First row: {}".format(tensor[0]))
    print("value of bottom right corner: {}".format(tensor[1, 1]))
    print("Last column: {}".format(tensor[:, -1]))
    print("First column: {}".format(tensor[..., 0]))


def base_calculate():
    x = Tensor(np.array([1, 2, 3]), mindspore.float32)
    y = Tensor(np.array([4, 5, 6]), mindspore.float32)

    # 加
    output_add = x + y
    # 减
    output_sub = x - y
    # 乘
    output_mul = x * y
    # 除
    output_div = y / x
    # 求余
    output_mod = y % x
    # 地板除
    output_floordiv = y // x

    print("add:", output_add)
    print("sub:", output_sub)
    print("mul:", output_mul)
    print("div:", output_div)
    print("mod:", output_mod)
    print("floordiv:", output_floordiv)


def show_concat(axis=0):
    data1 = Tensor(np.array([[0, 1], [2, 3]]).astype(np.float32))
    data2 = Tensor(np.array([[4, 5], [6, 7]]).astype(np.float32))
    output = ops.concat((data1, data2), axis=axis)
    print(output)
    print(output.shape, output.dtype)


def show_stack():
    data1 = Tensor(np.array([[0, 1], [2, 3]]).astype(np.float32))
    data2 = Tensor(np.array([[4, 5], [6, 7]]).astype(np.float32))
    output = ops.stack([data1, data2])

    print(output)
    print("shape:\n", output.shape)


def tensor2numpy():
    t = Tensor([1., 1., 1., 1., 1.])
    print(f"t: {t}", type(t))
    n = t.asnumpy()
    print(f"n: {n}", type(n))


def numpy2tensor():
    n = np.ones(5)
    t = Tensor.from_numpy(n)

    np.add(n, 1, out=n)
    print(f"n: {n}", type(n))
    print(f"t: {t}", type(t))


def show_CSRTensor():

    indptr = Tensor([0, 2, 5, 6])
    indices = Tensor([0, 3, 1, 2, 4, 2])
    values = Tensor([1., 2., 3., 4., 5., 6.], dtype=mindspore.float32)
    shape = (3, 5)
    # Make a CSRTensor
    csr_tensor = CSRTensor(indptr, indices, values, shape)

    print(csr_tensor.to_dense())
    print(csr_tensor.astype(mindspore.float64).dtype)


def show_COOTensor():
    indices = Tensor([[0, 1], [1, 2]], dtype=mindspore.int32)
    values = Tensor([1, 2], dtype=mindspore.float32)
    shape = (3, 4)

    # Make a COOTensor
    coo_tensor = COOTensor(indices, values, shape)
    print(coo_tensor.to_dense())
    print(coo_tensor.astype(mindspore.float64).dtype)


def show_RowTensor():
    indices = Tensor([0, 2])
    values = Tensor([[1, 2]], dtype=mindspore.float32)
    shape = (3, 2)

    row_tensor = RowTensor(indices, values, shape)
    print(row_tensor.dense_shape)
    # print(row_tensor.to_dense())


if __name__ =='__main__':
    # print("根据数据直接生成")
    # gen_tensor_base()
    #
    # print("从 NumPy 数组生成")
    # gen_tensor_numpy()
    #
    # print("使用 init 初始化器构造张量")
    # gen_tensor_init()
    #
    # print("继承另一个张量的属性，形成新的张量")
    # gen_tensor_extend()
    # show_tensor_index()

    # base_calculate()
    # print("axis=0")
    # show_concat(axis=0)
    # print()
    # print("axis=1")
    # show_concat(axis=1)

    # show_stack()
    # tensor2numpy()
    # numpy2tensor()
    # show_CSRTensor()
    # show_COOTensor()
    show_RowTensor()